package main

import (
	"flag"
	"fmt"
	"log"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"github.com/nsqio/go-nsq"
)

const (
	topicPrefix = "sub_bench_"
)

var (
	runfor     = flag.Duration("runfor", 10*time.Second, "duration of time to run")
	tcpAddress = flag.String("nsqd-tcp-address", "127.0.0.1:4150", "<addr>:<port> to connect to nsqd")
	topicCount = flag.Int("topic-count", 1, "topic count to receive messages on")
	size       = flag.Int("size", 200, "size of messages")
	batchSize  = flag.Int("batch-size", 200, "batch size of messages")
	deadline   = flag.String("deadline", "", "deadline to start the benchmark run")
)

var totalMsgCount int64

func main() {
	flag.Parse()
	if *topicCount < 1 {
		*topicCount = 1
	}
	var wg sync.WaitGroup

	log.SetPrefix("[bench_writer] ")

	msg := make([]byte, *size)
	batch := make([][]byte, *batchSize)
	for i := range batch {
		batch[i] = msg
	}

	goChan := make(chan int)
	rdyChan := make(chan int)
	workers := runtime.GOMAXPROCS(0)
	log.Printf("workers: %d", workers)
	for j := 0; j < workers; j++ {
		wg.Add(1)
		go func() {
			pubWorker(*runfor, *tcpAddress, *batchSize, batch, *topicCount, rdyChan, goChan)
			wg.Done()
		}()
		<-rdyChan
	}

	if *deadline != "" {
		t, err := time.Parse("2006-01-02 15:04:05", *deadline)
		if err != nil {
			log.Fatal(err)
		}
		d := t.Sub(time.Now())
		log.Printf("sleeping until %s (%s)", t, d)
		time.Sleep(d)
	}

	start := time.Now()
	close(goChan)
	wg.Wait()
	end := time.Now()
	duration := end.Sub(start)
	tmc := atomic.LoadInt64(&totalMsgCount)
	mbps := float64(tmc*int64(*size)) / duration.Seconds() / 1024 / 1024
	opspms := float64(tmc) / duration.Seconds() / 1000
	log.Printf("duration: %s - %.01fmb/s - %.01fops/ms - %.03fus/op",
		duration, mbps, opspms, float64(duration/time.Microsecond)/float64(tmc))
	log.Printf("Result: %d\t%.01f\t%.01f", *topicCount, mbps, opspms)
}

func pubWorker(td time.Duration, tcpAddr string, batchSize int, batch [][]byte, topicCount int, rdyChan chan int, goChan chan int) {
	config := nsq.NewConfig()
	config.DialTimeout = time.Second
	producer, err := nsq.NewProducer(tcpAddr, config)
	if err != nil {
		panic(err)
	}
	// Gracefully stop the producer.
	defer producer.Stop()
	producer.SetLogger(nil, 0)

	topics := make([]string, topicCount)
	for i := 0; i < topicCount; i++ {
		topics[i] = fmt.Sprintf("%s%d", topicPrefix, i+1)
	}

	rdyChan <- 1
	<-goChan
	var msgCount int64
	endTime := time.Now().Add(td)
	for {
		for i := 0; i < topicCount; i++ {
			if time.Now().After(endTime) {
				atomic.AddInt64(&totalMsgCount, msgCount)
				return
			}

			err = producer.MultiPublishAsync(topics[i], batch, nil)
			if err == nil {
				msgCount += int64(len(batch))
			}
			// ignore io timeout errors
		} // for i
	} // for
}
